from pocketflow import AsyncFlow, AsyncNode # 导入AsyncFlow
from nodes import InitialTriageNode, IntelligenceGatheringNode, ExpertJudgmentNode, GenerateReportNode, UnsupportedEventNode, RefineInvestigationNode

def create_soc_analysis_flow():
    # 实例化所有节点
    initial_triage_node = InitialTriageNode()
    intelligence_gathering_node = IntelligenceGatheringNode()
    expert_judgment_node = ExpertJudgmentNode()
    refine_investigation_node = RefineInvestigationNode() 
    
    # 报告生成节点，通过params区分报告类型
    true_positive_report_node = GenerateReportNode()
    false_positive_report_node = GenerateReportNode()
    uncertain_report_node = GenerateReportNode()
    
    # 处理不支持事件的节点
    unsupported_event_node = UnsupportedEventNode()

    # 定义流程连接
    initial_triage_node >> intelligence_gathering_node
    initial_triage_node - "unsupported_event" >> unsupported_event_node

    intelligence_gathering_node >> expert_judgment_node

    expert_judgment_node - "action_true_positive" >> true_positive_report_node
    expert_judgment_node - "action_false_positive" >> false_positive_report_node
    
    # 关键优化: 如果专家判断为“不确定”，则走向调查精炼节点，然后循环回情报收集
    expert_judgment_node - "action_uncertain" >> refine_investigation_node
    refine_investigation_node >> intelligence_gathering_node # 循环回情报收集，进行下一轮调查

    # 创建并返回AsyncFlow，指定起始节点
    return AsyncFlow(start=initial_triage_node)

if __name__ == "__main__":
    print("flow.py: 流程定义完成。请通过 main.py 运行测试。")
